import networkx as nx
from tools import generate_gas_network, visualize_network

def find_min_cut(G, leaks, sources, valves):
    """基于流向的有向图最小割算法"""
    H = G.copy()
    super_source = "super_source"
    super_sink = "super_sink"
    nodes = []
    nodes.append((super_source, {'type': 'super_source'}))
    nodes.append((super_sink, {'type': 'super_sink'}))
    H.add_nodes_from(nodes)

    for u in sources:
        H.add_edge(super_source, u, capacity=1e9)
    for u in leaks:
        H.add_edge(u, super_sink, capacity=1e9)

    for u, v in H.edges:
        if u not in valves:
            H[u][v]['capacity'] = 1e9  # 设置极大值        

    # 计算最小割
    min_cut_valve, partition = nx.minimum_cut(H, super_source, super_sink)
    reachable, non_reachable = partition
    print (min_cut_valve)
    print (partition)

    min_cut_set = []
    a = 0
    for u, v in H.edges():
        if u in reachable and v in non_reachable:
            if G.has_edge(u, v) and u in valves:
                min_cut_set.append((u, v))
                a += H[u][v]['capacity']
    print (min_cut_set)
    print (a)
    return min_cut_set

G, pos = generate_gas_network()
visualize_network(G, pos, "燃气管网可视化")
exit()
leaks = [n for n, attr in G.nodes(data=True) if attr['type'] == 'leak']
sources = [n for n, attr in G.nodes(data=True) if attr['type'] == 'source']
valves = [n for n, attr in G.nodes(data=True) if attr['type'] == 'valve']
print (valves)
min_cut_set = find_min_cut(G, leaks, sources, valves)
visualize_network(G, "燃气管网泄漏关阀优化方案", min_cut_set)

